import sqlite3
import numpy as np
import json

import torch
from PIL import Image
from scipy.spatial.distance import cosine
from AiUtil import extract_img_feature, load_features_from_hex


def load_features_from_db(db_path):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT word_id, word_feature,word_str FROM note_word")
    rows = cursor.fetchall()
    conn.close()

    features = {}
    strings = {}
    for row in rows:
        img_id = row[0]
        feature = load_features_from_hex(row[1])
        features[img_id] = feature
        strings[img_id] = row[2]

    return features, strings


def compare_features(feat1, feat2):
    return 1 - cosine(torch.flatten(feat1), torch.flatten(feat2))


def find_similar_images(features, query_feature, threshold=0.85):
    similar_images = []
    for img_id, feature in features.items():
        similarity = compare_features(query_feature, feature)
        if similarity > threshold:
            similar_images.append((img_id, similarity))
    return similar_images


if __name__ == '__main__':
    db_path = 'data/note.db'
    image_path = 'test/work1.png'

    features, strings = load_features_from_db(db_path)
    query_feature = extract_img_feature(Image.open(image_path))
    similar_images = find_similar_images(features, load_features_from_hex(query_feature), threshold=0.9)

    print("Similar images:")
    for img_id, similarity in similar_images:
        print(f"Image ID: {img_id},{strings[img_id]}, Similarity: {similarity}")
